#!/usr/bin/env python3

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

# 生成25个坐标点
trainData = np.random.randint(0, 100, (25, 2)).astype(np.float32)
# 25个0或1，用于类别区分
responses = np.random.randint(0, 2, (25, 1)).astype(np.float32)
# 取出红色点画三角形
red = trainData[responses.ravel() == 0]
plt.scatter(red[:, 0], red[:, 1], 80, 'r', '^')
# 取出蓝色点画方形
blue = trainData[responses.ravel() == 1]
plt.scatter(blue[:, 0], blue[:, 1], 80, 'b', 's')

# 添加新来者
newcomer = np.random.randint(0, 100, (1, 2)).astype(np.float32)
plt.scatter(newcomer[:, 0], newcomer[:, 1], 80, 'g', 'o')
# 创建kNN
knn = cv.ml.KNearest_create()
# 训练 k-NN 分类器
knn.train(trainData, cv.ml.ROW_SAMPLE, responses)
ret, results, neighbours, dist = knn.findNearest(newcomer, 3)
print("result: {}\n".format(results))
print("neighbours: {}\n".format(neighbours))
print("distance: {}\n".format(dist))
plt.show()
